#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
WebVNC 认证管理模块 - 新版本
支持管理员系统认证和协作者/观察者数据库认证
"""

import os
import sys
import sqlite3
import hashlib
import time
import datetime
from typing import Optional, Dict, List, Tuple

# 导入系统认证模块
try:
    from auth_os import authenticate
    SYSTEM_AUTH_AVAILABLE = True
except ImportError:
    SYSTEM_AUTH_AVAILABLE = False
    print("⚠️ 系统认证模块不可用，管理员认证功能将被禁用")

# 数据库文件路径
DB_PATH = os.path.join(os.path.dirname(__file__), "webvnc.db")

class User:
    """用户类"""
    def __init__(self, user_id: str, username: str, password: str, permission_level: int):
        self.user_id = user_id
        self.username = username
        self.password = password  # 管理员密码为空，其他用户为明文密码
        self.permission_level = permission_level
        self.login_time = None
        self.last_active = None

class AuthManager:
    """认证管理器"""
    
    # 权限级别定义
    PERMISSION_ADMIN = 4      # 管理员
    PERMISSION_COLLABORATOR = 3  # 协作者
    PERMISSION_OBSERVER = 2   # 观察者
    
    def __init__(self):
        self.init_database()
        self.sessions = {}  # session_id -> user_id
        self.users = {}     # user_id -> User
        self.kicked_sessions = set()
        # 启动定期清理过期会话的线程
        self.start_session_cleanup()
    
    def init_database(self):
        """初始化数据库"""
        conn = sqlite3.connect(DB_PATH)
        cursor = conn.cursor()
        
        # 创建用户表
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS users (
                user_id TEXT PRIMARY KEY,
                username TEXT UNIQUE NOT NULL,
                password TEXT NOT NULL,
                permission_level INTEGER NOT NULL,
                created_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        ''')
        
        # 创建日志表
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS logs (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                content TEXT NOT NULL
            )
        ''')
        
        conn.commit()
        conn.close()
    
    def is_initialized(self) -> bool:
        """检查系统是否已初始化"""
        conn = sqlite3.connect(DB_PATH)
        cursor = conn.cursor()
        
        # 检查是否只有一个管理员用户
        cursor.execute('''
            SELECT COUNT(*) FROM users 
            WHERE permission_level = ?
        ''', (self.PERMISSION_ADMIN,))
        
        admin_count = cursor.fetchone()[0]
        conn.close()
        
        return admin_count == 1
    
    def get_admin_user(self) -> Optional[User]:
        """获取管理员用户"""
        conn = sqlite3.connect(DB_PATH)
        cursor = conn.cursor()
        
        cursor.execute('''
            SELECT user_id, username, password, permission_level 
            FROM users 
            WHERE permission_level = ?
            LIMIT 1
        ''', (self.PERMISSION_ADMIN,))
        
        row = cursor.fetchone()
        conn.close()
        
        if row:
            return User(row[0], row[1], row[2], row[3])
        return None
    
    def reset_database(self):
        """重置数据库，用于重新初始化"""
        conn = sqlite3.connect(DB_PATH)
        cursor = conn.cursor()
        
        # 删除所有表
        cursor.execute('DROP TABLE IF EXISTS users')
        cursor.execute('DROP TABLE IF EXISTS logs')
        
        # 重新创建表
        self.init_database()
        
        self.log("系统数据库已重置，等待初始化")
        conn.close()
    
    def create_admin(self, username: str) -> bool:
        """创建管理员用户"""
        try:
            conn = sqlite3.connect(DB_PATH)
            cursor = conn.cursor()
            
            user_id = hashlib.md5(f"{username}_{time.time()}".encode()).hexdigest()
            
            cursor.execute('''
                INSERT INTO users (user_id, username, password, permission_level)
                VALUES (?, ?, ?, ?)
            ''', (user_id, username, "", self.PERMISSION_ADMIN))
            
            conn.commit()
            conn.close()
            
            self.log(f"管理员用户创建成功: {username}")
            return True
            
        except Exception as e:
            self.log(f"创建管理员失败: {e}")
            return False
    
    def add_user(self, username: str, password: str, permission_level: int, admin_user: str) -> Tuple[bool, str]:
        """添加用户（仅管理员）"""
        try:
            conn = sqlite3.connect(DB_PATH)
            cursor = conn.cursor()
            
            # 检查用户名是否已存在
            cursor.execute('SELECT COUNT(*) FROM users WHERE username = ?', (username,))
            if cursor.fetchone()[0] > 0:
                return False, "用户名已存在"
            
            # 检查权限级别是否有效
            if permission_level not in [self.PERMISSION_COLLABORATOR, self.PERMISSION_OBSERVER]:
                return False, "权限级别无效"
            
            user_id = hashlib.md5(f"{username}_{time.time()}".encode()).hexdigest()
            
            cursor.execute('''
                INSERT INTO users (user_id, username, password, permission_level)
                VALUES (?, ?, ?, ?)
            ''', (user_id, username, password, permission_level))
            
            conn.commit()
            conn.close()
            
            permission_name = self.get_permission_name(permission_level)
            self.log(f"管理员 {admin_user} 添加了{permission_name}用户: {username}")
            return True, "用户添加成功"
            
        except Exception as e:
            self.log(f"添加用户失败: {e}")
            return False, f"添加用户失败: {e}"
    
    def get_user_by_username(self, username: str) -> Optional[User]:
        """根据用户名获取用户"""
        conn = sqlite3.connect(DB_PATH)
        cursor = conn.cursor()
        
        cursor.execute('''
            SELECT user_id, username, password, permission_level 
            FROM users 
            WHERE username = ?
        ''', (username,))
        
        row = cursor.fetchone()
        conn.close()
        
        if row:
            return User(row[0], row[1], row[2], row[3])
        return None
    
    def verify_user(self, username: str, password: str, client_ip: str) -> Optional[User]:
        """验证用户"""
        user = self.get_user_by_username(username)
        if not user:
            return None
        
        # 管理员使用系统认证
        if user.permission_level == self.PERMISSION_ADMIN:
            if not SYSTEM_AUTH_AVAILABLE:
                self.log(f"系统认证不可用，管理员 {username} 登录失败")
                return None
            
            try:
                # 使用系统认证
                if authenticate(username, password):
                    self.log(f"管理员 {username} 从 {client_ip} 系统认证成功")
                    return user
                else:
                    self.log(f"管理员 {username} 从 {client_ip} 系统认证失败")
                    return None
            except Exception as e:
                self.log(f"管理员 {username} 系统认证错误: {e}")
                return None
        
        # 其他用户使用数据库认证
        else:
            if user.password == password:  # 明文密码比较
                permission_name = self.get_permission_name(user.permission_level)
                self.log(f"{permission_name} {username} 从 {client_ip} 数据库认证成功")
                return user
            else:
                permission_name = self.get_permission_name(user.permission_level)
                self.log(f"{permission_name} {username} 从 {client_ip} 数据库认证失败")
                return None
    
    def get_all_users(self) -> List[User]:
        """获取所有用户"""
        conn = sqlite3.connect(DB_PATH)
        cursor = conn.cursor()
        
        cursor.execute('''
            SELECT user_id, username, password, permission_level 
            FROM users 
            ORDER BY permission_level DESC, username ASC
        ''')
        
        users = []
        for row in cursor.fetchall():
            users.append(User(row[0], row[1], row[2], row[3]))
        
        conn.close()
        return users
    
    def delete_user(self, username: str, admin_user: str) -> Tuple[bool, str]:
        """删除用户（仅管理员）"""
        try:
            user = self.get_user_by_username(username)
            if not user:
                return False, "用户不存在"
            
            if user.permission_level == self.PERMISSION_ADMIN:
                return False, "不能删除管理员用户"
            
            conn = sqlite3.connect(DB_PATH)
            cursor = conn.cursor()
            
            cursor.execute('DELETE FROM users WHERE username = ?', (username,))
            
            conn.commit()
            conn.close()
            
            self.log(f"管理员 {admin_user} 删除了用户: {username}")
            return True, "用户删除成功"
            
        except Exception as e:
            self.log(f"删除用户失败: {e}")
            return False, f"删除用户失败: {e}"
    
    def create_session(self, user: User) -> str:
        """创建会话（单用户单会话）"""
        import uuid
        session_id = str(uuid.uuid4())
        
        # 踢出该用户的所有旧会话
        old_sessions = self.remove_user_sessions(user.user_id)
        
        # 如果有旧会话被踢出，记录到被踢出集合
        for old_session_id in old_sessions:
            self.kicked_sessions.add(old_session_id)
        
        # 创建新会话
        self.sessions[session_id] = user.user_id
        user.login_time = datetime.datetime.now()
        user.last_active = datetime.datetime.now()
        self.users[user.user_id] = user
        
        # 记录登录信息
        self.log(f"用户 {user.username} 创建新会话，踢出了 {len(old_sessions)} 个旧会话")
        
        return session_id
    
    def get_user_by_session(self, session_id: str) -> Optional[User]:
        """根据会话ID获取用户"""
        user_id = self.sessions.get(session_id)
        if user_id:
            user = self.users.get(user_id)
            if user:
                # 检查会话是否超时（30分钟）
                if user.last_active:
                    time_diff = datetime.datetime.now() - user.last_active
                    if time_diff.total_seconds() > 1800:  # 30分钟 = 1800秒
                        # 会话超时，清理会话
                        self.log(f"用户 {user.username} 会话超时，自动断开连接")
                        self.remove_session(session_id)
                        return None
                
                # 更新最后活动时间
                user.last_active = datetime.datetime.now()
                return user
        return None
    
    def remove_session(self, session_id: str):
        """移除会话"""
        if session_id in self.sessions:
            user_id = self.sessions[session_id]
            del self.sessions[session_id]
            if user_id in self.users:
                del self.users[user_id]
    
    def remove_user_sessions(self, user_id: str) -> List[str]:
        """移除用户的所有会话，返回被移除的会话ID列表"""
        sessions_to_remove = []
        for session_id, uid in self.sessions.items():
            if uid == user_id:
                sessions_to_remove.append(session_id)
        
        for session_id in sessions_to_remove:
            self.remove_session(session_id)
        
        return sessions_to_remove
    
    def get_all_active_users(self) -> Dict[str, User]:
        """获取所有活跃用户"""
        return self.users.copy()
    
    def is_session_kicked(self, session_id: str) -> bool:
        """检查会话是否被踢出"""
        return session_id in self.kicked_sessions
    
    def clear_kicked_session(self, session_id: str):
        """清除被踢出的会话标记"""
        self.kicked_sessions.discard(session_id)
    
    def get_permission_name(self, level: int) -> str:
        """获取权限级别名称"""
        names = {
            self.PERMISSION_ADMIN: "管理员",
            self.PERMISSION_COLLABORATOR: "协作者",
            self.PERMISSION_OBSERVER: "观察者"
        }
        return names.get(level, "未知")
    
    def get_user_permissions(self, user: User) -> Dict[str, bool]:
        """获取用户权限"""
        permissions = {
            "can_view_desktop": True,  # 所有用户都可以查看桌面
            "can_control_desktop": user.permission_level >= self.PERMISSION_COLLABORATOR,
            "can_use_chat": True,  # 所有用户都可以使用聊天
            "can_view_users": user.permission_level >= self.PERMISSION_COLLABORATOR,
            "can_manage_users": user.permission_level >= self.PERMISSION_ADMIN,
            "can_kick_users": user.permission_level >= self.PERMISSION_ADMIN,
            "can_add_users": user.permission_level >= self.PERMISSION_ADMIN,
            "can_delete_users": user.permission_level >= self.PERMISSION_ADMIN,
            "can_view_logs": user.permission_level >= self.PERMISSION_ADMIN,
        }
        return permissions
    
    def log(self, content: str):
        """记录日志"""
        try:
            conn = sqlite3.connect(DB_PATH)
            cursor = conn.cursor()
            
            cursor.execute('''
                INSERT INTO logs (content) VALUES (?)
            ''', (content,))
            
            # 保持最多1000条日志
            cursor.execute('''
                DELETE FROM logs 
                WHERE id NOT IN (
                    SELECT id FROM logs 
                    ORDER BY timestamp DESC 
                    LIMIT 1000
                )
            ''')
            
            conn.commit()
            conn.close()
            
            print(f"📝 {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')} - {content}")
            
        except Exception as e:
            print(f"❌ 日志记录失败: {e}")
    
    def get_recent_logs(self, limit: int = 50) -> List[Tuple[str, str]]:
        """获取最近的日志"""
        try:
            conn = sqlite3.connect(DB_PATH)
            cursor = conn.cursor()
            
            cursor.execute('''
                SELECT timestamp, content 
                FROM logs 
                ORDER BY timestamp DESC 
                LIMIT ?
            ''', (limit,))
            
            logs = cursor.fetchall()
            conn.close()
            
            return logs
            
        except Exception as e:
            print(f"❌ 获取日志失败: {e}")
            return []
    
    def cleanup_expired_sessions(self):
        """清理过期的会话"""
        current_time = datetime.datetime.now()
        expired_sessions = []
        
        for session_id, user_id in list(self.sessions.items()):
            user = self.users.get(user_id)
            if user and user.last_active:
                time_diff = current_time - user.last_active
                if time_diff.total_seconds() > 1800:  # 30分钟
                    expired_sessions.append(session_id)
        
        # 清理过期会话
        for session_id in expired_sessions:
            user_id = self.sessions.get(session_id)
            user = self.users.get(user_id)
            if user:
                self.log(f"清理过期会话: 用户 {user.username}")
            self.remove_session(session_id)
        
        return len(expired_sessions)
    
    def start_session_cleanup(self):
        """启动定期清理过期会话的线程"""
        import threading
        
        def cleanup_worker():
            while True:
                try:
                    # 每5分钟检查一次过期会话
                    time.sleep(300)  # 5分钟 = 300秒
                    expired_count = self.cleanup_expired_sessions()
                    if expired_count > 0:
                        print(f"🧹 清理了 {expired_count} 个过期会话")
                except Exception as e:
                    print(f"❌ 清理过期会话时出错: {e}")
        
        cleanup_thread = threading.Thread(target=cleanup_worker, daemon=True)
        cleanup_thread.start()
        print("🚀 会话清理线程已启动")

# 全局认证管理器实例
auth_manager = AuthManager()
